#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Copyright 2020 Huawei Technologies Co., Ltd
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import acl
import sys
import struct
import numpy as np
sys.path.append("..")
from resource.resource import ascend_resource
from common.log import Log, check_ret
from common.const import ACL_MEMCPY_HOST_TO_HOST, \
                     ACL_MEMCPY_HOST_TO_DEVICE, \
                     ACL_MEMCPY_DEVICE_TO_HOST, \
                     ACL_MEMCPY_DEVICE_TO_DEVICE, \
                     ACL_MEM_MALLOC_HUGE_FIRST, \
                     ACL_MEM_MALLOC_HUGE_ONLY, \
                     ACL_MEM_MALLOC_NORMAL_ONLY, \
                     PIXEL_FORMAT_YUV_SEMIPLANAR_420, \
                     PIXEL_FORMAT_YVU_SEMIPLANAR_420, \
                     PIXEL_FORMAT_YUV_SEMIPLANAR_422, \
                     PIXEL_FORMAT_YVU_SEMIPLANAR_422, \
                     PIXEL_FORMAT_YUV_SEMIPLANAR_444, \
                     PIXEL_FORMAT_YVU_SEMIPLANAR_444, \
                     DVPP, DEVICE, HOST, \
                     DEBUG,INFO,WARNING,ERROR, \
                     STATIC_MODE, DYNAMIC_BATCH, DYNAMIC_RESOLUTION, \
                     NPY_BYTE

class AscendModel():
    def __init__(self,
                 context,
                 stream,
                 model_path):
        self.class_name = self.__class__.__name__
        self.context = context  # pointer
        self.stream = stream

        self.model_desc = None  # pointer when using
        self.model_path = model_path  # string
        self.model_id = None  # pointer
        self.model_input_ptr = None
        self.model_output_ptr = None
        self.model_input_size = None
        self.model_output_size = None
        self.model_batch = None
        self.input_dataset = None
        self.output_dataset = None
        self.host_output = []
        self.output_info = []
   
        self.memory = ascend_resource.memory_pool       

        self.init_resource()
        self.is_create = True

    def init_resource(self):
        #acl.rt.set_context(self.context)
        Log.do_log(INFO, self.class_name, 'init_resource', ' init resource stage start')
        self.gen_model_desc()
        self.init_dataset()
        Log.do_log(INFO, self.class_name, 'init_resource', ' init resource stage success')

    def init_dataset(self):
        self.input_dataset = acl.mdl.create_dataset()
        self.output_dataset = acl.mdl.create_dataset()
    
    def gen_model_desc(self):
        self.model_id, ret = acl.mdl.load_from_file(self.model_path)
        check_ret(self.class_name, '_gen_model_desc', 'fail to load mode from file', ret)
        self.model_desc = acl.mdl.create_desc()
        ret = acl.mdl.get_desc(self.model_desc, self.model_id)
        check_ret(self.class_name, '_gen_model_desc', 'fail to get model id', ret)    

    def gen_model_data_num(self, opt):
        if opt == 'input':
            size = acl.mdl.get_num_inputs(self.model_desc)
        elif opt == 'output':
            size = acl.mdl.get_num_outputs(self.model_desc)
        else:
            check_ret(self.class_name, 'gen_model_data_num', 'opt is not supported', 1)
        return size

    def gen_datainfo(self,opt):
        size = self.gen_model_data_num(opt)
        ptr_list = [0] * size
        size_list = [0] * size
        batch_list = [0]* size
        for i in range(0, size):
            if opt == 'input':
                size_list[i] = acl.mdl.get_input_size_by_index(self.model_desc, i)
                ptr_list[i] = self.memory.memory_allocation_(size_list[i], DEVICE)
                dims, ret = acl.mdl.get_input_dims(self.model_desc,i)
                check_ret(self.class_name, 'send_data_to_model', 'fail to get_input_dims', ret)
                batch_list[i] = dims["dims"][0]
            if opt == 'output':
                size_list[i] = acl.mdl.get_output_size_by_index(self.model_desc, i)
                ptr_list[i] = self.memory.memory_allocation_(size_list[i], DEVICE)
                dims, ret = acl.mdl.get_output_dims(self.model_desc,i)
                check_ret(self.class_name, 'send_data_to_model', 'fail to get_output_dims', ret)
                batch_list[i] = dims["dims"][0]
                datatype = acl.mdl.get_output_data_type(self.model_desc, i)
                self.output_info.append({"shape": tuple(dims["dims"]), "type": datatype})

        Log.do_log(INFO, self.class_name, 'gen_model_data', 'init model {} success'.format(opt))
        return ptr_list, size_list, batch_list 
 
    def gen_databuffer(self, opt, buffer_ptr:list, buffer_size:list):
        if opt == 'input':
            dataset = self.input_dataset
        if opt == 'output':
            dataset = self.output_dataset

        for i in range(0, len(buffer_ptr)):
            temp_buffer = acl.create_data_buffer(buffer_ptr[i], buffer_size[i])
            _, ret = acl.mdl.add_dataset_buffer(dataset, temp_buffer)
            try :
                check_ret(self.class_name, '_gen_databuffer', 'fail to create {} dataset buffer'.format(opt),ret)
            except :
                acl.destroy_data_buffer(dataset)
                return
        Log.do_log(INFO,self.class_name, '_gen_databuffer', 'create model data buffer success')

    def gen_dataset(self, policy, width, height, batch):
        if policy == STATIC_MODE:
           self.policy = policy
           if self.is_create == True:
                self.is_create = False
                #input
                self.model_input_ptr, self.model_input_size, self.model_batch = self.gen_datainfo('input')
                self.gen_databuffer('input', self.model_input_ptr, self.model_input_size)
                #output
                self.model_output_ptr, self.model_output_size, _ = self.gen_datainfo('output')
                self.gen_databuffer('output', self.model_output_ptr, self.model_output_size)
        else:
            check_ret(self.class_name, 'gen_dataset', 'policy is not supported', 1)

    def run(self, policy, dvpp_buffer: list, dvpp_size: list, width = None, height = None, batch = None):
        self.gen_dataset(policy, width, height, batch)
        self.send_data_to_model(dvpp_buffer, dvpp_size)      
        self.forward()
        self.send_result_to_host()
        return self.host_output

    def forward(self):
        ret = acl.mdl.execute_async(self.model_id,
                              self.input_dataset,
                              self.output_dataset,
                              self.stream)
        check_ret(self.class_name, 'forward', 'fail to send desc to model', ret)
        ret = acl.rt.synchronize_stream(self.stream)
        check_ret(self.class_name, 'forward', 'fail to execute model', ret)
        Log.do_log(INFO, self.class_name, 'forward', 'model execute stage success')

    def send_data_to_model(self, ptr:list, size:list):
        num = self.gen_model_data_num('input')
        if len(ptr) != len(size) or num != len(ptr):
           check_ret(self.class_name, 'send_data_to_model', 'input data num do not match model requirement', 2)
        for i in range(0, num):
            if size[i] == self.model_input_size[i]:                   
                ret = acl.rt.memcpy(self.model_input_ptr[i], size[i], ptr[i], size[i], ACL_MEMCPY_DEVICE_TO_DEVICE)
                check_ret(self.class_name, 'send_data_to_model', 'fail to memcpy for ptr[{}]'.format(i), ret)
            else:
                check_ret(self.class_name, 'send_data_to_model', 'input data size do not match model requirement', 2)
        Log.do_log(INFO, self.class_name, 'send_data_to_model', 'send data to model success')

    def send_result_to_host(self):
        self.host_output = []
        num = acl.mdl.get_dataset_num_buffers(self.output_dataset)
        for i in range(num):
            temp_output_buffer = acl.mdl.get_dataset_buffer(self.output_dataset, i)
            infer_output_ptr = acl.get_data_buffer_addr(temp_output_buffer)
            infer_output_size = acl.get_data_buffer_size(temp_output_buffer)
            output_host = self.memory.memory_allocation_(infer_output_size, HOST)
            ret = acl.rt.memcpy(output_host, infer_output_size, infer_output_ptr,
                                infer_output_size, ACL_MEMCPY_DEVICE_TO_HOST)
            check_ret(self.class_name, 'send_result_to_host', 'fail to memcy to host', ret)
            self.host_output.append([output_host, infer_output_size])
        Log.do_log(INFO, self.class_name, 'send_result_to_host', 'send data to host success')

    def _release_databuffer(self,opt):
        if opt == 'input':
            dataset = self.input_dataset
        if opt == 'output':
            dataset = self.output_dataset

        for i in range(0, acl.mdl.get_dataset_num_buffers(dataset)):
            data_buf = acl.mdl.get_dataset_buffer(dataset, i)
            if data_buf:
                ret = acl.destroy_data_buffer(data_buf)
                check_ret(self.class_name, '_release_dataset', 'fail to destory {} databuffer'.format(opt), ret)

    def _release_dataset(self):
        if self.input_dataset:
            ret = acl.mdl.destroy_dataset(self.input_dataset)
            check_ret(self.class_name, '_release_dataset', 'fail to destory input dataset', ret)
        if self.output_dataset:
            ret = acl.mdl.destroy_dataset(self.output_dataset)
            check_ret(self.class_name, '_release_dataset', 'fail to destory output dataset', ret)

    def release_model_resource(self, context_id):
        if context_id is not None:
            ret = acl.rt.set_context(context_id)
        self._release_databuffer('input')
        self._release_databuffer('output')
        self._release_dataset()
        if self.model_id:
            ret = acl.mdl.unload(self.model_id)
            check_ret(self.class_name, 'release_model_resource', 'fail to unload mode', ret)
        if self.model_desc:
            ret = acl.mdl.destroy_desc(self.model_desc)
            check_ret(self.class_name, 'release_model_resource', 'fail to destory desc', ret)
        Log.do_log(INFO, self.class_name, 'release_model_resource','release source success')

    
